import os
import sys
import yaml
import logging
import asyncio
from typing import List, Dict, Any
from pathlib import Path

# 添加项目根目录到Python路径
project_root = str(Path(__file__).parent)
sys.path.append(project_root)

from graph_rag.store import Neo4jGraphStore
from graph_rag.embedder import GraphEmbedder
from graph_rag.searcher import GraphSearcher
from graph_rag.reasoner import GraphAugmenter, GraphReasoner


class GraphRAG:
    """GraphRAG系统"""
    
    def __init__(self, config_path: str):
        self.config = self._load_config(config_path)
        self.logger = self._setup_logging()
        
        # 初始化组件
        self.store = Neo4jGraphStore()
        self.embedder = GraphEmbedder()
        self.searcher = GraphSearcher()
        self.augmenter = GraphAugmenter()
        self.reasoner = GraphReasoner()
        
        # 初始化所有组件
        self._initialize_components()
        
    def _load_config(self, config_path: str) -> Dict[str, Any]:
        """加载配置文件"""
        with open(config_path, 'r', encoding='utf-8') as f:
            return yaml.safe_load(f)
            
    def _setup_logging(self) -> logging.Logger:
        """设置日志"""
        logging_config = self.config["logging"]
        
        # 创建日志目录
        log_file = logging_config["file"]
        os.makedirs(os.path.dirname(log_file), exist_ok=True)
        
        # 配置日志
        logging.basicConfig(
            level=logging_config["level"],
            format=logging_config["format"],
            handlers=[
                logging.FileHandler(log_file),
                logging.StreamHandler()
            ]
        )
        
        return logging.getLogger(__name__)
        
    def _initialize_components(self) -> None:
        """初始化所有组件"""
        try:
            # 初始化存储
            self.store.initialize({
                "uri": self.config["neo4j"]["uri"],
                "username": self.config["neo4j"]["username"],
                "password": self.config["neo4j"]["password"]
            })
            
            # 初始化嵌入器
            self.embedder.initialize({
                "model_name": self.config["embedder"]["model_name"]
            })
            
            # 初始化搜索器
            self.searcher.initialize({
                "store": {
                    "class": Neo4jGraphStore,
                    **self.config["neo4j"]
                },
                "embedder": {
                    "class": GraphEmbedder,
                    **self.config["embedder"]
                },
                **self.config["search"]
            })
            
            # 初始化增强器
            self.augmenter.initialize({
                "model_name": self.config["llm"]["model_name"],
                "temperature": self.config["llm"]["temperature"],
                "openai_api_key": os.getenv("OPENAI_API_KEY"),
                **self.config["augmentation"]
            })
            
            # 初始化推理器
            self.reasoner.initialize({
                "model_name": self.config["llm"]["model_name"],
                "temperature": self.config["llm"]["temperature"],
                "openai_api_key": os.getenv("OPENAI_API_KEY"),
                **self.config["reasoning"]
            })
            
            self.logger.info("所有组件初始化完成")
            
        except Exception as e:
            self.logger.error(f"组件初始化失败: {str(e)}")
            raise
            
    async def query(self, query: str) -> Dict[str, Any]:
        """处理查询"""
        try:
            # 1. 图搜索
            self.logger.info(f"执行图搜索: {query}")
            search_results = await self.searcher.search(
                query,
                node_types=self.config["search"]["node_types"],
                max_results=self.config["search"]["max_results"],
                subgraph_depth=self.config["search"]["subgraph_depth"]
            )
            
            if not search_results:
                return {
                    "answer": "未找到相关信息",
                    "confidence": 0.0,
                    "evidence": []
                }
            
            # 2. 图增强
            self.logger.info("执行图增强")
            augmented_subgraphs = []
            if self.config["augmentation"]["enabled"]:
                for result in search_results:
                    augmented = await self.augmenter.augment(query, result.subgraph)
                    augmented_subgraphs.append(augmented)
            else:
                augmented_subgraphs = [result.subgraph for result in search_results]
            
            # 3. 图推理
            self.logger.info("执行图推理")
            reasoning_results = []
            for subgraph in augmented_subgraphs:
                result = await self.reasoner.reason(query, subgraph)
                if result["confidence"] >= self.config["reasoning"]["confidence_threshold"]:
                    reasoning_results.append(result)
            
            if not reasoning_results:
                return {
                    "answer": "无法得出可靠结论",
                    "confidence": 0.0,
                    "evidence": []
                }
            
            # 4. 合并结果
            final_result = self._merge_reasoning_results(reasoning_results)
            return final_result
            
        except Exception as e:
            self.logger.error(f"查询处理失败: {str(e)}")
            return {
                "answer": "处理查询时出错",
                "confidence": 0.0,
                "evidence": [],
                "error": str(e)
            }
            
    def _merge_reasoning_results(self, results: List[Dict[str, Any]]) -> Dict[str, Any]:
        """合并多个推理结果"""
        if not results:
            return {
                "answer": "无推理结果",
                "confidence": 0.0,
                "evidence": []
            }
            
        # 按置信度排序
        sorted_results = sorted(results, key=lambda x: x["confidence"], reverse=True)
        best_result = sorted_results[0]
        
        # 收集所有证据
        all_evidence = []
        for result in sorted_results:
            all_evidence.extend(result["evidence"])
        
        # 限制证据数量
        max_evidence = self.config["reasoning"]["evidence_count"]
        all_evidence = list(dict.fromkeys(all_evidence))[:max_evidence]
        
        return {
            "answer": best_result["answer"],
            "confidence": best_result["confidence"],
            "evidence": all_evidence,
            "reasoning_steps": best_result["reasoning_steps"],
            "uncertainty": best_result["uncertainty"]
        }


async def main():
    # 检查环境变量
    if not os.getenv("OPENAI_API_KEY"):
        print("错误: 未设置OPENAI_API_KEY环境变量")
        return
        
    # 创建GraphRAG实例
    config_path = os.path.join(project_root, "config", "config.yaml")
    rag = GraphRAG(config_path)
    
    # 示例查询
    queries = [
        "什么是向量检索？",
        "图神经网络的应用场景有哪些？",
        "知识图谱和传统数据库的区别是什么？",
        "如何评估RAG系统的性能？"
    ]
    
    # 处理查询
    for query in queries:
        print(f"\n处理查询: {query}")
        print("-" * 50)
        
        result = await rag.query(query)
        
        print("\n答案:", result["answer"])
        print("\n置信度:", result["confidence"])
        print("\n证据:")
        for evidence in result["evidence"]:
            print(f"- {evidence}")
        print("\n推理步骤:")
        for step in result["reasoning_steps"]:
            print(f"- {step}")
        print("\n不确定性:", result["uncertainty"])
        print("=" * 50)


if __name__ == "__main__":
    asyncio.run(main()) 